package com.lagou.edu.gateway.filter;

import lombok.Data;
import org.apache.commons.lang.time.DateUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.stream.Collectors;

/**
 * ip防刷
 *
 * @author liangzj
 * @date 2021/6/27 15:23
 */
@RefreshScope
@Component
public class IpDefenseFilter implements GlobalFilter, Ordered {

    @Value("${defense.cycle.second}")
    private Integer second;
    @Value("${defense.frequency.max-allow-time}")
    private Integer maxAllowTime;

    /**
     * 统计所有Map<ip, statistics>
     */
    private static final ConcurrentHashMap<String, List<RequestStatisticsDTO>> ipRequestStatisticsMap = new ConcurrentHashMap();

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        Date now = new Date();
        // 先记录本次请求
        addRequestRecord(request, now);
        // 删除second秒前数据
        Date startTime = DateUtils.addSeconds(now, -second);
        Map<String, List<RequestStatisticsDTO>> cleanedStatisticsMap = ipRequestStatisticsMap
                .entrySet()
                .parallelStream()
                .peek(record -> record.getValue().removeIf(statistics -> statistics.getLastTime().before(startTime)))
                .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
        // 获取second秒内此ip对所有接口的所有请求
        List<RequestStatisticsDTO> ipRequestRecords = cleanedStatisticsMap.get(request.getRemoteAddress().getAddress().getHostAddress());
        // 没有请求过直接放行
        if (CollectionUtils.isEmpty(ipRequestRecords)) {
            return chain.filter(exchange);
        }
        // 超过允许次数直接拒绝
        long timeCount = ipRequestRecords.stream().filter(record -> request.getURI().getPath().equals(record.getUri())).count();
        if (timeCount >= maxAllowTime) {
            ServerHttpResponse response = exchange.getResponse();
            response.setStatusCode(HttpStatus.SEE_OTHER);
            return response.setComplete();
        }
        return chain.filter(exchange);
    }

    private void addRequestRecord(ServerHttpRequest request, Date requestTime) {
        String remoteAddress = request.getRemoteAddress().getAddress().getHostAddress();
        RequestStatisticsDTO statisticsDTO = new RequestStatisticsDTO(request.getURI().getPath(), requestTime);

        List<RequestStatisticsDTO> requestStatisticsDTOS = ipRequestStatisticsMap.get(remoteAddress);
        if (requestStatisticsDTOS == null) {
            ipRequestStatisticsMap.put(remoteAddress, new CopyOnWriteArrayList<>());
        }

        ipRequestStatisticsMap.get(remoteAddress).add(statisticsDTO);
    }

    @Override
    public int getOrder() {
        return 1;
    }

    @Data
    private static class RequestStatisticsDTO {
        /**
         * 请求链接
         */
        private String uri;

        /**
         * 最后请求时间
         */
        private Date lastTime;

        public RequestStatisticsDTO(String uri, Date lastTime) {
            this.uri = uri;
            this.lastTime = lastTime;
        }
    }
}
